from io import BytesIO
import json
import os
import tempfile

from app.pkg.paillier_zk import PaillierZK



class PaillierZKFileHandler:
    def __init__(self, zk, oss_handler,zk_class=PaillierZK):
        self.zk = zk
        self.oss_handler = oss_handler  # 引入OSS处理器
        self.zk_class = zk_class

    def encrypt_file(self, file_obj, filename, chunk_size=512):
        print("🔐 开始文件加密...")

        content = file_obj.read()
        print(f"📦 文件读取完成，总字节数: {len(content)}")

        encrypted_data = []
        block_index = 0
        for i in range(0, len(content), chunk_size):
            chunk = content[i:i + chunk_size]
            for byte in chunk:
                c, proof = self.zk.encrypt_block(byte)
                encrypted_data.append({
                    'cipher': c,
                    'proof': proof
                })
                block_index += 1
                if block_index % 100 == 0:
                    print(f"✅ 已加密 {block_index} 字节")

        encrypted_file_data = json.dumps(encrypted_data)
        print(f"✅ 加密完成，总加密块数: {block_index}")

        # 将加密数据保存到内存中
        encrypted_file_obj = BytesIO(encrypted_file_data.encode('utf-8'))

        # 上传到OSS
        self.oss_handler.upload_file(encrypted_file_obj, f'paillier/{os.path.splitext(filename)[0]}_encrypted.json')
        # 2. 上传密钥
        key_data = {
            "N": str(self.zk.N),
            "N_sq": str(self.zk.N_sq),
            "lambda": str(self.zk.lambda_),
            "mu": str(self.zk.mu),
            "g": str(self.zk.g),
        }
        key_json_bytes = BytesIO(json.dumps(key_data).encode('utf-8'))
        key_filename = f"paillier/{os.path.splitext(filename)[0]}_key.json"
        self.oss_handler.upload_file(key_json_bytes, key_filename)

    # def decrypt_file(self, filename, output_path):
    def decrypt_file(self, filename):
        print("🔓 开始从 OSS 下载并解密文件...")
        
        try:
            original_name = os.path.splitext(filename)[0]
            key_filename = f"paillier/{original_name}_key.json"
            encrypted_filename = f"paillier/{original_name}_encrypted.json"

            temp_dir = tempfile.mkdtemp()
            key_file = os.path.join(temp_dir, f"{original_name}_key.json")
            data_file = os.path.join(temp_dir, f"{original_name}_encrypted.json")

            # 下载
            self.oss_handler.download_file(key_filename, key_file)
            self.oss_handler.download_file(encrypted_filename, data_file)
            with open(key_file, 'r') as f:
                key_data = json.load(f)
                self.zk = self.zk_class()
                self.zk.N = int(key_data['N'])
                self.zk.N_sq = int(key_data['N_sq'])
                self.zk.lambda_ = int(key_data['lambda'])
                self.zk.mu = int(key_data['mu'])
                self.zk.g = int(key_data['g'])
            with open(data_file, 'r') as f:
                encrypted_data = json.load(f)
                print(f"📂 加载加密数据完成，块总数: {len(encrypted_data)}")

            # 从 OSS 下载加密文件，假设 OSS 返回文件内容为字节流
            # encrypted_data_bytes = self.oss_handler.download_file_stream(encrypted_filename)

            # 将字节流转换为 JSON 格式
            # encrypted_data = json.loads(encrypted_data_bytes.decode('utf-8'))

            result = []
            for idx, item in enumerate(encrypted_data):
                c = item['cipher']
                proof = item['proof']
                if not self.zk.verify_block(c, proof):
                    raise ValueError(f"❌ 第 {idx} 块 零知识证明验证失败")
                m = self.zk.decrypt_block(c)
                result.append(m)
                if (idx + 1) % 100 == 0:
                    print(f"✅ 已验证并解密 {idx + 1} 块")

            decrypted_file_data = bytes(result)

            # 将解密后的数据保存到文件
            # with open(output_path, 'wb') as f:
                # f.write(decrypted_file_data)
            # print(f"✅ 解密完成，文件保存到: {output_path}")
            print(f"✅ 解密完成")
            return decrypted_file_data
        except Exception as e:
            print(f"❌ 解密过程中发生错误: {e}")
